"""
    
    File: util.py
    Description:
        
        Class to limit connection to a given service.
    
    Author: Kyle Vogt
    Date  : January 18th, 2008
    Copyright (c) 2008, Justin.tv, Inc.

"""

from twisted.python import log
import time
import sys
try:
    from sizer import scanner
    from sizer import formatting
except:
    log.msg('PySizer not installed, will not be able to profile memory.')
from twisted.internet import reactor, defer

class ConnectionManager:
    
    def __init__(self, ip_limit = 50, limit = 5000):
        "Allow override of connections per ip and total connections"
        self.clients = []
        self.hosts = {}
        self.ip_connection_limit = ip_limit
        self.connection_limit = limit
    
    def add(self, client):
        "Add a new client if it doesn't exceed connection limits"
        conns = self.hosts.get(client.ip, 0)
        if conns >= self.ip_connection_limit:
            log.msg("Connection rejected, at ip connection limit of %s" % self.ip_connection_limit)
            return False
        elif len(self.clients) + 1 > self.connection_limit:
            log.msg("Connection rejected, at global connection limit of %s" % self.connection_limit)
            return False
        else:
            self.hosts[client.ip] = conns + 1
            self.clients.append(client)
            log.msg("Connection accepted. %s has %s connections. %s total" % (client.ip, self.hosts[client.ip], len(self.clients)))
            return True
    
    def remove(self, client):
        if client in self.clients:
            conns = self.hosts.get(client.ip, 0)
            if conns > 0: conns -= 1
            self.hosts[client.ip] = conns
            self.clients.remove(client)
            return True
        else: return False

class Profiler:
    
    def __init__(self):
        self.time_start = {}
        self.time_used = {}
        self.time_length = 10.0
    
    def benchmark(self, func, *kw, **kwargs):
        if func not in self.time_start:
            self.reset(func)
        
        st = time.time()
        func(*kw, **kwargs)
        self.time_used[func] += (time.time() - st)
        if (time.time() - self.time_start[func]) > self.time_length:
            log.msg('Used %.3fs out of 10s for func %s' % (self.time_used[func], func))
            self.reset(func)
    
    def reset(self, func):
        self.time_start[func] = time.time()
        self.time_used[func] = 0
        
class MemoryProfiler:
        
    def sample(self):
        def profile():
            self.objs = scanner.Objects()
        return reactor.callLater(0.0, profile)
        
